import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from enable.api import ComponentEditor

class SaturatedPressure(HasTraits):
    Temp = Range(low=0.1, high=100.0, value=60.0)
    data = Property(Array, observe=['Temp'])
    plot = Instance(Plot)
    def _plot_default(self):
        T = np.arange(0, self.Temp, 0.1)  # 温度
        E = np.exp(53.67957 - 6743.769 / (T + 273.15) - 4.8451 * np.log(T + 273.15))  # 气压 hpa
        plot = Plot(ArrayPlotData(T=T, E=E))
        plot.title = 'Saturated Pressure vs Temperature'
        plot.x_axis.title = "Temperature /C"
        plot.y_axis.title = "Saturated Pressure / hPa"
        plot.plot(("T", "E"), name='Saturated pressure', color=palette11[0])
        return plot

    def _get_data(self):
        data={}
        T = np.arange(0, self.Temp, 0.1)  # 温度
        E = np.exp(53.67957 - 6743.769 / (T + 273.15) - 4.8451 * np.log(T + 273.15))  # 气压 hpa
        data['T']=T
        data['E'] = E
        return data

    @observe("data")
    def _update_plot(self, event):
        self.plot.data.update_data(
            T=self.data['T'],
            E=self.data['E'],
        )
    traits_view = View(
        Group(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            )
        ),
        Item('Temp', label='最高温度/℃', editor_args={'mode': 'slider'}),
    )

class GasPressure(HasTraits):
    Height = Range(low=100, high=100000.0, value=10000.0)
    data = Property(Array, observe=['Height'])
    plot = Instance(Plot)
    def _plot_default(self):
        Height = np.arange(0, self.Height, 1)  # 高度 m
        P = 1013.25*(1-2.256e-5*Height)**5.256  # 气压 hpa
        plot = Plot(ArrayPlotData(H=Height, P=P))
        plot.title = 'Gas Pressure vs Height'
        plot.x_axis.title = "Height /m"
        plot.y_axis.title = "Gas Pressure / hPa"
        plot.plot(("H", "P"), name='Gas pressure', color=palette11[0])
        return plot

    def _get_data(self):
        data={}
        Height = np.arange(0, self.Height, 1)  # 高度 m
        P = 1013.25 * (1 - 2.256e-5 * Height) ** 5.256  # 气压 hpa
        data['H']=Height
        data['P'] = P
        return data

    @observe("data")
    def _update_plot(self, event):
        self.plot.data.update_data(
            H=self.data['H'],
            P=self.data['P'],
        )
    traits_view = View(
        Group(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            )
        ),
        Item('Height', label='最高高度/m', editor_args={'mode': 'slider'}),
    )
class Curvatureradius(HasTraits):
    latitude = Range(low=0.0, high=90.0, value=90)
    data = Property(Array, observe=['latitude'])
    plot = Instance(Plot)
    def _plot_default(self):
        latitude = np.arange(0.1, self.latitude+0.1, 0.1)  # 纬度
        a = 6378160.
        b = 6356775.
        theta = np.where(latitude == 90.0, np.pi / 2 - 1e-12, latitude / 180 * np.pi)
        a1 = (b ** 2/(a**2) + a ** 2/(b**2) * (np.tan(theta)) ** 2)
        a2 = (1/(a ** 2) + 1/(b ** 2) * (np.tan(theta)) ** 2)
        rho = 1/(a*b)*(a1/a2)**(3/2)
        np.min(rho)
        plot = Plot(ArrayPlotData(L=latitude, R=rho))
        plot.title = 'Curvature radius vs latitude'
        plot.x_axis.title = "latitude /degree"
        plot.y_axis.title = "Curvature radius"
        plot.plot(("L", "R"), name='Curvature radius', color=palette11[0])
        return plot

    def _get_data(self):
        data = {}
        latitude = np.arange(0.1, self.latitude-0.1, 0.1)  # 纬度
        a = 6378160.
        b = 6356775.
        theta = np.where(latitude == 90.0, np.pi / 2 - 1e-12, latitude / 180 * np.pi)
        a1 = (b ** 2 / (a ** 2) + a ** 2 / (b ** 2) * (np.tan(theta)) ** 2)
        a2 = (1 / (a ** 2) + 1 / (b ** 2) * (np.tan(theta)) ** 2)
        rho = 1 / (a * b) * (a1 / a2) ** (3 / 2)
        data['L']= latitude
        data['R'] = rho
        return data

    @observe("data")
    def _update_plot(self, event):
        self.plot.data.update_data(
            L=self.data['L'],
            R=self.data['R'],
        )

    traits_view = View(
        Group(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            )
        ),
        Item('latitude', label='纬度/°', editor_args={'mode': 'slider'}),
    )

class GasN(HasTraits):
    Height = Range(low=100, high=100000.0, value=10000.0)
    humidity= Range(low=0.0, high=100.0, value=80.0)
    Temp = Range(low=0.1, high=100.0, value=60.0)
    data = Property(Array, observe=['Height','humidity','Temp'])
    plot = Instance(Plot)
    def _plot_default(self):
        Height = np.arange(0, self.Height, 1)  # 高度 m
        E = np.exp(53.67957 - 6743.769 / (self.Temp + 273.15) - 4.8451 * np.log(self.Temp + 273.15))
        P = 1013.25*(1-2.256e-5*Height)**5.256
        N=77.6/(self.Temp+273.15)*(P+4810*self.humidity/100*E /(self.Temp+273.15))
        plot = Plot(ArrayPlotData(H=Height, N=N))
        plot.title = 'Atmospheric refract N vs Height'
        plot.x_axis.title = "Height /m"
        plot.y_axis.title = "Atmospheric refract N"
        plot.plot(("H", "N"), name='Atmospheric refract', color=palette11[0])
        return plot

    def _get_data(self):
        data={}
        Height = np.arange(0, self.Height, 1)  # 高度 m
        E = np.exp(53.67957 - 6743.769 / (self.Temp + 273.15) - 4.8451 * np.log(self.Temp + 273.15))
        P = 1013.25 * (1 - 2.256e-5 * Height) ** 5.256
        N = 77.6 / (self.Temp + 273.15) * (P + 4810 * self.humidity / 100 * E / (self.Temp + 273.15))
        data['H']=Height
        data['N'] = N
        return data

    @observe("data")
    def _update_plot(self, event):
        self.plot.data.update_data(
            H=self.data['H'],
            N=self.data['N'],
        )
    traits_view = View(
        Group(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            )
        ),
        Item('Height', label='最高高度/m', editor_args={'mode': 'slider'}),
        Item('humidity', label='湿度', editor_args={'mode': 'slider'}),
        Item('Temp', label='温度/℃', editor_args={'mode': 'slider'}),
    )
class Form(HasTraits):
    SaturatedPressure = Instance(SaturatedPressure, ())
    GasPressure = Instance(GasPressure, ())
    Curvatureradius=Instance(Curvatureradius, ())
    GasN=Instance(GasN, ())
    traits_view = View(
        Tabbed(
            UItem('SaturatedPressure', style='custom', label='饱和水汽压'),
            UItem('GasPressure', style='custom', label='大气压强'),
            UItem('Curvatureradius', style='custom', label='曲率半径与纬度'),
            UItem('GasN', style='custom', label='大气折射率N'),
        ),
        title='大气传播',
        width=1200,
        height=800,
        resizable=True,
    )

if __name__ == '__main__':
    form = Form()
    form.configure_traits()
